"""通过对数据集中所有的bounding boxes进行聚类，获取yolo的base anchors。
"""
import argparse
import numpy as np
from mmcv import Config

from mmdet.datasets import build_dataset
from mmhead import datasets  # noqa: F401,F403


def _iou(box: np.array, clusters: np.array) -> np.array:
    """Calculates the Intersection over Union (IoU) between a box and k clusters.

    Args:
        box (np.array): ground truth bounding box width and height
        clusters (np.array): clusters, shape is [k, 2],
            k is the number of clusters.

    Returns:
        np.array: iou between the box and clusters, shape is (k, )
    """
    x = np.minimum(clusters[:, 0], box[0])
    y = np.minimum(clusters[:, 1], box[1])
    if np.count_nonzero(x == 0) > 0 or np.count_nonzero(y == 0) > 0:
        raise ValueError("Box has no area")

    intersection = x * y
    box_area = box[0] * box[1]
    cluster_area = clusters[:, 0] * clusters[:, 1]
    iou = intersection / (box_area + cluster_area - intersection)
    return iou


def _kmeans(boxes: np.array, k: int, dist=np.median) -> np.array:
    """Calculates k-means clustering with the Intersection over Union (IoU) metric.

    Args:
        boxes (np.array): all the ground truth bounding boxes,
            shape is [r, 2], r is the number of ground truth bounding boxes,
            2 means width and height.
        k (int): number of clusters.
        dist ([type], optional): [description]. Defaults to np.median.

    Returns:
        np.array: clusters, shape is (k, 2)
    """
    rows = boxes.shape[0]
    distances = np.empty((rows, k))
    last_mean_iou = 0.
    np.random.seed()
    # the Forgy method will fail if the whole array contains the same rows
    clusters = boxes[np.random.choice(rows, k, replace=False)]
    for iter in range(10000):  # 最多迭代10000次
        for row in range(rows):
            distances[row] = 1 - _iou(boxes[row], clusters)
        nearest_clusters = np.argmin(distances, axis=1)

        d = sum([distances[i, j] for i, j in enumerate(nearest_clusters)])
        mean_iou = 1 - d / len(nearest_clusters)
        if abs(mean_iou - last_mean_iou) <= 0.01:
            break
        last_mean_iou = mean_iou
        if iter % 10 == 0:
            print(F"iter: {iter}, mean iou is: {mean_iou:.2f}")

        for cluster in range(k):
            clusters[cluster] = dist(
                boxes[nearest_clusters == cluster], axis=0)
    print(F"last mean iou is: {last_mean_iou:.2f}")
    clusters = clusters.tolist()
    clusters = sorted(clusters, key=lambda x: x[0] * x[1])
    clusters = [[int(c[0]), int(c[1])] for c in clusters]
    return clusters


def argument_parser():
    parser = argparse.ArgumentParser(description="dataset test")
    parser.add_argument(
        "--config-file",
        default="",
        metavar="FILE",
        help="path to config file")
    args = parser.parse_args()
    return args


def get_dataset_boxes(dataset) -> np.array:
    gt_bboxes_list = list()
    for i, data in enumerate(dataset):
        gt_bboxes = data["gt_bboxes"].data.numpy()
        if len(gt_bboxes) == 0:
            continue
        for gt_bbox in gt_bboxes:
            w = gt_bbox[2] - gt_bbox[0]
            h = gt_bbox[3] - gt_bbox[1]
            if w == 0 or h == 0:
                print(F"data {i} has zero area boxes...")
        gt_bboxes_list.append(gt_bboxes)
    gt_bboxes = np.concatenate(gt_bboxes_list, axis=0)
    boxes = np.empty(shape=(gt_bboxes.shape[0], 2), dtype=gt_bboxes.dtype)
    boxes[:, 0] = gt_bboxes[:, 2] - gt_bboxes[:, 0]
    boxes[:, 1] = gt_bboxes[:, 3] - gt_bboxes[:, 1]
    return boxes


def main():
    args = argument_parser()
    config_file = args.config_file
    cfg = Config.fromfile(config_file)
    dataset = build_dataset(cfg.data.train)
    boxes = get_dataset_boxes(dataset)
    print(F"sample number: {len(dataset)}, box number: {len(boxes)}")

    print("begin to cluster...")
    clusters = _kmeans(boxes, k=9)
    print(F"clusters is:\n\t{clusters}")


if __name__ == "__main__":
    main()
